import paramiko
import re

class SSHCollector:
    def __init__(self, host):
        self.host = host
        
    def collect_all(self):
        """收集所有SSH监控数据"""
        try:
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(
                self.host.ip_address,
                port=self.host.ssh_port,
                username=self.host.ssh_username,
                password=self.host.ssh_password
            )

            # 收集CPU使用率
            self.host.cpu_usage = self._get_cpu_usage(ssh)
            
            # 收集内存使用率
            self.host.memory_usage = self._get_memory_usage(ssh)
            
            # 收集磁盘使用率
            self.host.disk_usage = self._get_disk_usage(ssh)
            
            ssh.close()
            self.host.save()
            
        except Exception as e:
            print(f"SSH数据收集失败: {str(e)}")
            raise

    def _get_cpu_usage(self, ssh):
        stdin, stdout, stderr = ssh.exec_command("top -bn1 | grep 'Cpu(s)'")
        cpu_output = stdout.read().decode()
        cpu_usage = float(re.search(r'(\d+\.\d+) id', cpu_output).group(1))
        return 100 - cpu_usage

    def _get_memory_usage(self, ssh):
        stdin, stdout, stderr = ssh.exec_command("free | grep Mem")
        memory_output = stdout.read().decode()
        total, used = map(int, re.findall(r'Mem:\s+(\d+)\s+(\d+)', memory_output)[0])
        return (used / total) * 100

    def _get_disk_usage(self, ssh):
        stdin, stdout, stderr = ssh.exec_command("df / | tail -1")
        disk_output = stdout.read().decode()
        usage = int(re.search(r'(\d+)%', disk_output).group(1))
        return usage 